{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "adf859aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training class distribution:\n",
      "label\n",
      "NORM     7364\n",
      "OTHER    6792\n",
      "RBBB     1237\n",
      "AFIB      977\n",
      "1dAVb     638\n",
      "LBBB      408\n",
      "AFLT       53\n",
      "Name: count, dtype: int64\n",
      "\n",
      "Test class distribution:\n",
      "label\n",
      "NORM     1841\n",
      "OTHER    1699\n",
      "RBBB      309\n",
      "AFIB      245\n",
      "1dAVb     159\n",
      "LBBB      102\n",
      "AFLT       13\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import ast\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# Load metadata\n",
    "df = pd.read_csv('C:\\\\Users\\\\mstew\\\\OneDrive\\\\Curtin Univeristy\\\\COMP6011\\\\Task 3\\\\ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.1\\\\ptbxl_database.csv')\n",
    "\n",
    "# Parse scp_codes column to dict\n",
    "df['scp_codes'] = df['scp_codes'].apply(ast.literal_eval)\n",
    "\n",
    "# Map all diagnostic labels to SCP codes\n",
    "priority_map = {\n",
    "    '1dAVb': ['1AVB'],  # First-degree AV block\n",
    "    'RBBB': ['IRBBB', 'CRBBB'],  # Right bundle branch block (partial + complete)\n",
    "    'LBBB': ['ILBBB', 'CLBBB'],  # Left bundle branch block (partial + complete)\n",
    "    'AFLT': ['AFLT'],  # Atrial flutter\n",
    "    'AFIB': ['AFIB'],  # Atrial fibrillation\n",
    "    'NORM': ['NORM'],  # Normal\n",
    "}\n",
    "\n",
    "# Function to assign single label based on priority\n",
    "def assign_priority_label(scp_dict):\n",
    "    for label, codes in priority_map.items():\n",
    "        if any(code in scp_dict for code in codes):\n",
    "            return label\n",
    "    return 'OTHER'\n",
    "\n",
    "# Assign single label\n",
    "df['label'] = df['scp_codes'].apply(assign_priority_label)\n",
    "\n",
    "# Filter to only those samples with labels in the desired classes\n",
    "df = df[df['label'].isin(['1dAVb', 'RBBB', 'LBBB', 'AFLT', 'AFIB', 'OTHER', 'NORM'])]\n",
    "\n",
    "# Split into training and test sets\n",
    "train_df, test_df = train_test_split(df, test_size=0.2, stratify=df['label'], random_state=42)\n",
    "\n",
    "# Optional: display class balance\n",
    "print(\"Training class distribution:\")\n",
    "print(train_df['label'].value_counts())\n",
    "print(\"\\nTest class distribution:\")\n",
    "print(test_df['label'].value_counts())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "83e765b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 17469/17469 [02:41<00:00, 108.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved to test_data_balanced.npz with shape (17469, 1000, 12)\n",
      "Shape of X: (17469, 1000, 12)\n",
      "Shape of y: (17469,)\n",
      "First few labels: ['RBBB' 'AFIB' 'AFIB' 'OTHER' 'OTHER']\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import wfdb\n",
    "from tqdm import tqdm\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "# Assuming you've already created train_df or test_df with a 'label' column\n",
    "# and that 'filename_lr' gives paths to 100 Hz signals\n",
    "\n",
    "def load_and_save_ecg_data(df, data_dir, output_filename):\n",
    "    signals = []\n",
    "    labels = []\n",
    "\n",
    "    for _, row in tqdm(df.iterrows(), total=len(df)):\n",
    "        path = f\"{data_dir}/{row['filename_lr']}\"\n",
    "        try:\n",
    "            record = wfdb.rdrecord(path)\n",
    "            sig = record.p_signal  # shape: (n_samples, 12)\n",
    "            signals.append(sig)\n",
    "            labels.append(row['label'])\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading {path}: {e}\")\n",
    "\n",
    "    # Ensure all samples are the same length (optional: pad/truncate if needed)\n",
    "    min_length = min([s.shape[0] for s in signals])\n",
    "    signals = [s[:min_length] for s in signals]  # truncate all to min length\n",
    "\n",
    "    # Convert to arrays\n",
    "    X = np.array(signals)  # shape: (num_samples, min_length, 12)\n",
    "    y = np.array(labels)\n",
    "\n",
    "    \n",
    "\n",
    "    # Encode labels\n",
    "    label_encoder = LabelEncoder()\n",
    "    y_encoded = label_encoder.fit_transform(y)\n",
    "\n",
    "    # Save to .npz\n",
    "    np.savez(output_filename, X=X, y=y_encoded, label_names=label_encoder.classes_)\n",
    "\n",
    "    print(f\"Saved to {output_filename} with shape {X.shape}\")\n",
    "    return X, y\n",
    "\n",
    "# Load and save the ECG data for the test set\n",
    "X, y = load_and_save_ecg_data(train_df, data_dir='C:\\\\Users\\\\mstew\\\\OneDrive\\\\Curtin Univeristy\\\\COMP6011\\\\Task 3\\\\ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.1', output_filename='test_data_balanced.npz')\n",
    "\n",
    "# look at the shape of the data\n",
    "print(f\"Shape of X: {X.shape}\")\n",
    "print(f\"Shape of y: {y.shape}\")\n",
    "# Check the first few labels\n",
    "print(\"First few labels:\", y[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0630e4ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "def jitter(signal, sigma=0.01):\n",
    "    return signal + np.random.normal(0, sigma, size=signal.shape)\n",
    "\n",
    "def scaling(signal, sigma=0.1):\n",
    "    factor = np.random.normal(loc=1.0, scale=sigma)\n",
    "    return signal * factor\n",
    "\n",
    "def time_stretch(signal, stretch_factor=0.1):\n",
    "    orig_len = signal.shape[1]\n",
    "    factor = 1 + np.random.uniform(-stretch_factor, stretch_factor)\n",
    "    new_len = int(orig_len * factor)\n",
    "    x = np.linspace(0, 1, orig_len)\n",
    "    f = interp1d(x, signal, kind='linear', axis=1, fill_value='extrapolate')\n",
    "    x_new = np.linspace(0, 1, new_len)\n",
    "    stretched = f(x_new)\n",
    "    if stretched.shape[1] > orig_len:\n",
    "        return stretched[:, :orig_len]\n",
    "    else:\n",
    "        # pad if needed\n",
    "        pad_width = orig_len - stretched.shape[1]\n",
    "        return np.pad(stretched, ((0, 0), (0, pad_width)), mode='edge')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "5575f714",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from sklearn.utils import shuffle\n",
    "\n",
    "def generate_augmented_samples(X_class, num_needed, label):\n",
    "    num_synth = int(num_needed // 3)\n",
    "    num_aug = num_needed - num_synth\n",
    "    \n",
    "    # Synthetic oversampling\n",
    "    synth_samples = X_class[np.random.randint(0, len(X_class), num_synth)]\n",
    "\n",
    "    # Augmentation: jitter, scaling, stretch (round robin)\n",
    "    aug_samples = []\n",
    "    funcs = [jitter, scaling, time_stretch]\n",
    "    for i in range(num_aug):\n",
    "        idx = np.random.randint(0, len(X_class))\n",
    "        func = funcs[i % len(funcs)]\n",
    "        aug = func(X_class[idx])\n",
    "        aug_samples.append(aug)\n",
    "\n",
    "    aug_samples = np.array(aug_samples)\n",
    "    X_boosted = np.concatenate([synth_samples, aug_samples], axis=0)\n",
    "    y_boosted = np.full(len(X_boosted), label)\n",
    "    return X_boosted, y_boosted\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "860be2e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(17469,)\n",
      "1237\n",
      "977\n",
      "6792\n",
      "7364\n",
      "638\n",
      "408\n",
      "53\n",
      "Training class distribution after balancing:\n",
      "{'AFIB': 2000, 'NORM': 7364, 'OTHER': 6792, 'LBBB': 2000, 'AFLT': 2000, 'RBBB': 2000, '1dAVb': 2000}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Split the data into features and labels\n",
    "label_col = 'label'\n",
    "X_train = train_df.drop(columns=[label_col])\n",
    "y_train = train_df[label_col]\n",
    "\n",
    "print(y_train.shape)\n",
    "\n",
    "target_min = 2000\n",
    "X_train_boosted = []\n",
    "y_train_boosted = []\n",
    "\n",
    "class_counts = Counter(y_train)\n",
    "\n",
    "import wfdb\n",
    "\n",
    "data_dir = 'C:\\\\Users\\\\mstew\\\\OneDrive\\\\Curtin Univeristy\\\\COMP6011\\\\Task 3\\\\ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.1'\n",
    "\n",
    "# For each class, load signals and augment if needed\n",
    "for cls, count in class_counts.items():\n",
    "    print(count)\n",
    "    X_cls_df = X_train[y_train == cls]\n",
    "    y_cls = y_train[y_train == cls]\n",
    "    # Load signals for this class\n",
    "    signals = []\n",
    "    for fname in X_cls_df['filename_lr']:\n",
    "        path = f\"{data_dir}/{fname}\"\n",
    "        try:\n",
    "            record = wfdb.rdrecord(path)\n",
    "            sig = record.p_signal.T  # shape: (12, n_samples)\n",
    "            signals.append(sig)\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading {path}: {e}\")\n",
    "    # Make all signals the same length\n",
    "    min_length = min([s.shape[1] for s in signals])\n",
    "    signals = [s[:, :min_length] for s in signals]\n",
    "    X_cls_signals = np.stack(signals, axis=0)  # shape: (num_samples, 12, min_length)\n",
    "    y_cls_arr = np.array([cls] * len(X_cls_signals))\n",
    "    X_train_boosted.append(X_cls_signals)\n",
    "    y_train_boosted.append(y_cls_arr)\n",
    "    if count < target_min:\n",
    "        needed = target_min - count\n",
    "        X_new, y_new = generate_augmented_samples(X_cls_signals, needed, cls)\n",
    "        X_train_boosted.append(X_new)\n",
    "        y_train_boosted.append(y_new)\n",
    "\n",
    "X_train_balanced = np.concatenate(X_train_boosted, axis=0)\n",
    "y_train_balanced = np.concatenate(y_train_boosted, axis=0)\n",
    "\n",
    "X_train_balanced, y_train_balanced = shuffle(X_train_balanced, y_train_balanced, random_state=42)\n",
    "\n",
    "print(\"Training class distribution after balancing:\")\n",
    "print(dict(Counter(y_train_balanced)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "38c767b3",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'test_df' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[15], line 9\u001b[0m\n\u001b[0;32m      6\u001b[0m X_test_boosted \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m      7\u001b[0m y_test_boosted \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m----> 9\u001b[0m test_counts \u001b[38;5;241m=\u001b[39m Counter(\u001b[43mtest_df\u001b[49m[label_col])\n\u001b[0;32m     11\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m \u001b[38;5;28mcls\u001b[39m, count \u001b[38;5;129;01min\u001b[39;00m test_counts\u001b[38;5;241m.\u001b[39mitems():\n\u001b[0;32m     12\u001b[0m     X_cls_df \u001b[38;5;241m=\u001b[39m test_df[test_df[label_col] \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mcls\u001b[39m]\n",
      "\u001b[1;31mNameError\u001b[0m: name 'test_df' is not defined"
     ]
    }
   ],
   "source": [
    "# Duplicate underrepresented test samples by loading ECG signals and duplicating as needed\n",
    "label_col = 'label'\n",
    "target_test_min = 400\n",
    "data_dir = 'C:\\\\Users\\\\mstew\\\\OneDrive\\\\Curtin Univeristy\\\\COMP6011\\\\Task 3\\\\ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.1'\n",
    "\n",
    "X_test_boosted = []\n",
    "y_test_boosted = []\n",
    "\n",
    "test_counts = Counter(test_df[label_col])\n",
    "\n",
    "for cls, count in test_counts.items():\n",
    "    X_cls_df = test_df[test_df[label_col] == cls]\n",
    "    signals = []\n",
    "    for fname in X_cls_df['filename_lr']:\n",
    "        path = f\"{data_dir}/{fname}\"\n",
    "        try:\n",
    "            record = wfdb.rdrecord(path)\n",
    "            sig = record.p_signal.T  # shape: (12, n_samples)\n",
    "            signals.append(sig)\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading {path}: {e}\")\n",
    "    if len(signals) == 0:\n",
    "        continue\n",
    "    min_length = min([s.shape[1] for s in signals])\n",
    "    signals = [s[:, :min_length] for s in signals]\n",
    "    X_cls_signals = np.stack(signals, axis=0)  # (num_samples, 12, min_length)\n",
    "    y_cls_arr = np.array([cls] * len(X_cls_signals))\n",
    "    X_test_boosted.append(X_cls_signals)\n",
    "    y_test_boosted.append(y_cls_arr)\n",
    "    if count < target_test_min:\n",
    "        needed = target_test_min - count\n",
    "        reps = needed // len(X_cls_signals) + 1\n",
    "        X_dup = np.tile(X_cls_signals, (reps, 1, 1))[:needed]\n",
    "        y_dup = np.full(len(X_dup), cls)\n",
    "        X_test_boosted.append(X_dup)\n",
    "        y_test_boosted.append(y_dup)\n",
    "\n",
    "X_test_balanced = np.concatenate(X_test_boosted, axis=0)\n",
    "y_test_balanced = np.concatenate(y_test_boosted, axis=0)\n",
    "\n",
    "X_test_balanced, y_test_balanced = shuffle(X_test_balanced, y_test_balanced, random_state=42)\n",
    "\n",
    "print(\"Test class distribution after balancing:\")\n",
    "print(dict(Counter(y_test_balanced)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "886ecd00",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez('X_train_balanced.npz', X=X_train_balanced)\n",
    "np.savez('y_train_balanced.npz', y=y_train_balanced)\n",
    "np.savez('X_test_balanced.npz', X=X_test_balanced)\n",
    "np.savez('y_test_balanced.npz', y=y_test_balanced)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d3688f17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved to train_data_balanced.npz with shape (24156, 12, 1000)\n",
      "Shape of X: (24156, 12, 1000)\n",
      "Shape of y: (24156,)\n",
      "First few labels: ['AFIB' 'NORM' 'OTHER' 'OTHER' 'OTHER']\n"
     ]
    }
   ],
   "source": [
    "# Use the already prepared balanced arrays\n",
    "X = X_train_balanced\n",
    "y = y_train_balanced\n",
    "\n",
    "# Encode labels\n",
    "label_encoder = LabelEncoder()\n",
    "y_encoded = label_encoder.fit_transform(y)\n",
    "\n",
    "# Save to .npz\n",
    "output_filename = 'train_data_balanced.npz'\n",
    "np.savez(output_filename, X=X, y=y_encoded, label_names=label_encoder.classes_)\n",
    "\n",
    "print(f\"Saved to {output_filename} with shape {X.shape}\")\n",
    "print(f\"Shape of X: {X.shape}\")\n",
    "print(f\"Shape of y: {y.shape}\")\n",
    "print(\"First few labels:\", y[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f114dab9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved to train_data.npz with shape (17469, 1000, 12)\n"
     ]
    }
   ],
   "source": [
    "output_filename = 'train_data.npz'\n",
    "# Encode labels\n",
    "label_encoder = LabelEncoder()\n",
    "y_encoded = label_encoder.fit_transform(y)\n",
    "\n",
    "# Save to .npz\n",
    "np.savez(output_filename, X=X, y=y_encoded, label_names=label_encoder.classes_)\n",
    "\n",
    "print(f\"Saved to {output_filename} with shape {X.shape}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34062749",
   "metadata": {},
   "source": [
    "The below section is for prepping the MIT-BIH arrhythmia database for use with our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f08e5eb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wfdb\n",
    "import numpy as np\n",
    "import os\n",
    "from scipy.signal import resample\n",
    "\n",
    "def load_mitbih_data(path, target_length=1000, target_fs=100):\n",
    "    \"\"\"\n",
    "    Load and prepare ECG data from MIT-BIH Arrhythmia Database.\n",
    "\n",
    "    Args:\n",
    "        path (str): Directory containing .dat, .atr, and .hea files.\n",
    "        target_length (int): Number of samples expected by model (e.g., 1000).\n",
    "        target_fs (int): Target sampling rate for model (e.g., 100 Hz).\n",
    "\n",
    "    Returns:\n",
    "        list of np.ndarray: ECGs with shape (2, 1000, 1) from MLII and V1 leads.\n",
    "        list of str: Corresponding rhythm annotations per sample.\n",
    "    \"\"\"\n",
    "    ecg_segments = []\n",
    "    annotations = []\n",
    "\n",
    "    files = [f[:-4] for f in os.listdir(path) if f.endswith(\".hea\")]\n",
    "    print(f\"Found {len(files)} records.\")\n",
    "\n",
    "    for record_name in files:\n",
    "        try:\n",
    "            # Read signal and annotations\n",
    "            record = wfdb.rdrecord(os.path.join(path, record_name))\n",
    "            ann = wfdb.rdann(os.path.join(path, record_name), 'atr')\n",
    "            signal = record.p_signal  # shape (N, 2) usually: MLII, V1\n",
    "            original_fs = record.fs\n",
    "\n",
    "            # Leads present\n",
    "            leads = record.sig_name\n",
    "            lead_indices = {name: i for i, name in enumerate(leads)}\n",
    "            if 'MLII' not in lead_indices or 'V1' not in lead_indices:\n",
    "                continue\n",
    "\n",
    "            # Extract and resample the leads\n",
    "            ml_ii = signal[:, lead_indices['MLII']]\n",
    "            v1 = signal[:, lead_indices['V1']]\n",
    "\n",
    "            # Resample to 100 Hz\n",
    "            new_len = int(len(ml_ii) * target_fs / original_fs)\n",
    "            ml_ii_resampled = resample(ml_ii, new_len)\n",
    "            v1_resampled = resample(v1, new_len)\n",
    "\n",
    "            # Sliding window extraction\n",
    "            step = target_length\n",
    "            for i in range(0, len(ml_ii_resampled) - target_length, step):\n",
    "                segment = np.zeros((12, target_length))  # Fill all 12 leads with zero\n",
    "                segment[0] = ml_ii_resampled[i:i+target_length]  # Simulate as lead I\n",
    "                segment[6] = v1_resampled[i:i+target_length]     # Simulate as V1\n",
    "\n",
    "                ecg_segments.append(segment[..., np.newaxis])  # shape: (12, 1000, 1)\n",
    "                annotations.append(\"UNKNOWN\")  # Replace later with label if needed\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Failed to process {record_name}: {e}\")\n",
    "\n",
    "    print(f\"Prepared {len(ecg_segments)} samples.\")\n",
    "    return ecg_segments, annotations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d5ceae5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mitbih_classes(path):\n",
    "    from collections import Counter\n",
    "    all_classes = Counter()\n",
    "\n",
    "    files = [f[:-4] for f in os.listdir(path) if f.endswith(\".hea\")]\n",
    "    for record_name in files:\n",
    "        try:\n",
    "            ann = wfdb.rdann(os.path.join(path, record_name), 'atr')\n",
    "            symbols = ann.symbol\n",
    "            all_classes.update(symbols)\n",
    "        except Exception as e:\n",
    "            print(f\"Could not read annotations for {record_name}: {e}\")\n",
    "\n",
    "    print(\"Annotation symbols and counts:\")\n",
    "    for symbol, count in all_classes.items():\n",
    "        print(f\"{symbol}: {count}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bdc33517",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Annotation symbols and counts:\n",
      "+: 1291\n",
      "N: 75052\n",
      "A: 2546\n",
      "V: 7130\n",
      "~: 616\n",
      "|: 132\n",
      "Q: 33\n",
      "/: 7028\n",
      "f: 982\n",
      "x: 193\n",
      "F: 803\n",
      "j: 229\n",
      "L: 8075\n",
      "a: 150\n",
      "J: 83\n",
      "R: 7259\n",
      "[: 6\n",
      "!: 472\n",
      "]: 6\n",
      "E: 106\n",
      "S: 2\n",
      "\": 437\n",
      "e: 16\n"
     ]
    }
   ],
   "source": [
    "# Replace with your MIT-BIH dataset directory\n",
    "mitbih_path = \"C:\\\\Users\\\\mstew\\\\OneDrive\\\\Curtin Univeristy\\\\COMP6011\\\\Task 3\\\\mit-bih-arrhythmia-database-1.0.0\"\n",
    "\n",
    "# Print class labels\n",
    "get_mitbih_classes(mitbih_path)\n",
    "\n",
    "# # Load data for model\n",
    "# X_test_mitbih, y_labels = load_mitbih_data(mitbih_path)\n",
    "\n",
    "# # Convert to numpy array\n",
    "# X_test_mitbih = np.array(X_test_mitbih)  # shape: (N, 12, 1000, 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "62c4cb9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wfdb\n",
    "import numpy as np\n",
    "from collections import defaultdict\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "import os\n",
    "\n",
    "# Fixed label order for model alignment\n",
    "FIXED_LABEL_ORDER = ['1dAVb', 'AFIB', 'AFLT', 'LBBB', 'NORM', 'OTHER', 'RBBB']\n",
    "\n",
    "# Mapping annotation keywords to class names\n",
    "MITBIH_TO_CUSTOM = {\n",
    "    '1dAVb': ['1AV'],\n",
    "    'RBBB': ['R'],\n",
    "    'LBBB': ['L'],\n",
    "    'AFLT': ['AFL'],\n",
    "    'AFIB': ['AFIB'],\n",
    "    'NORM': ['N']\n",
    "}\n",
    "\n",
    "def map_annotation(comment):\n",
    "    for cls, terms in MITBIH_TO_CUSTOM.items():\n",
    "        for term in terms:\n",
    "            if term in comment:\n",
    "                return cls\n",
    "    return \"OTHER\"\n",
    "\n",
    "def prepare_mitbih_dataset(data_dir, record_list, max_per_class=1500):\n",
    "    class_counts = defaultdict(int)\n",
    "    X_data = []\n",
    "    y_data = []\n",
    "\n",
    "    for rec in record_list:\n",
    "        record_path = os.path.join(data_dir, rec)\n",
    "        try:\n",
    "            record = wfdb.rdrecord(record_path)\n",
    "            annotation = wfdb.rdann(record_path, 'atr')\n",
    "        except Exception as e:\n",
    "            print(f\"Skipping {rec} due to error: {e}\")\n",
    "            continue\n",
    "\n",
    "        signals = record.p_signal.T  # Shape: (n_leads, n_samples)\n",
    "        fs = record.fs\n",
    "        if signals.shape[0] != 2:\n",
    "            print(f\"Skipping {rec}: expected 2 leads, found {signals.shape[0]}\")\n",
    "            continue\n",
    "\n",
    "        for i, (sample, sym) in enumerate(zip(annotation.sample, annotation.aux_note)):\n",
    "            label = map_annotation(sym)\n",
    "\n",
    "            if class_counts[label] >= max_per_class:\n",
    "                continue\n",
    "\n",
    "            # 1000-sample segment centered on beat\n",
    "            start = sample - 500\n",
    "            end = sample + 500\n",
    "\n",
    "            if start < 0 or end > signals.shape[1]:\n",
    "                continue\n",
    "\n",
    "            segment = signals[:, start:end]  # shape: (2, 1000)\n",
    "\n",
    "            if segment.shape[1] != 1000:\n",
    "                continue\n",
    "\n",
    "            # Expand to (12, 1000) by repeating leads if necessary\n",
    "            padded_segment = np.zeros((12, 1000))\n",
    "            padded_segment[:2, :] = segment\n",
    "\n",
    "            X_data.append(padded_segment[..., np.newaxis])  # shape: (12, 1000, 1)\n",
    "            y_data.append(label)\n",
    "            class_counts[label] += 1\n",
    "\n",
    "\n",
    "    # Encode y_data with fixed label order\n",
    "    label_to_index = {label: idx for idx, label in enumerate(FIXED_LABEL_ORDER)}\n",
    "    y_encoded = np.array([label_to_index[label] for label in y_data])\n",
    "\n",
    "    print(\"Class distribution:\", dict(class_counts))\n",
    "    print(\"Label mapping:\", label_to_index)\n",
    "\n",
    "    X = np.array(X_data)\n",
    "    y = y_encoded\n",
    "\n",
    "    return X, y, label_to_index\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "60460005",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping 110 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/110.hea'\n",
      "Skipping 120 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/120.hea'\n",
      "Skipping 125 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/125.hea'\n",
      "Skipping 126 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/126.hea'\n",
      "Skipping 127 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/127.hea'\n",
      "Skipping 128 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/128.hea'\n",
      "Skipping 129 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/129.hea'\n",
      "Skipping 130 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/130.hea'\n",
      "Skipping 131 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/131.hea'\n",
      "Skipping 132 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/132.hea'\n",
      "Skipping 133 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/133.hea'\n",
      "Skipping 134 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/134.hea'\n",
      "Skipping 135 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/135.hea'\n",
      "Skipping 136 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/136.hea'\n",
      "Skipping 137 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/137.hea'\n",
      "Skipping 138 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/138.hea'\n",
      "Skipping 139 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/139.hea'\n",
      "Skipping 140 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/140.hea'\n",
      "Skipping 141 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/141.hea'\n",
      "Skipping 142 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/142.hea'\n",
      "Skipping 143 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/143.hea'\n",
      "Skipping 144 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/144.hea'\n",
      "Skipping 145 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/145.hea'\n",
      "Skipping 146 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/146.hea'\n",
      "Skipping 147 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/147.hea'\n",
      "Skipping 148 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/148.hea'\n",
      "Skipping 149 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/149.hea'\n",
      "Skipping 150 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/150.hea'\n",
      "Skipping 151 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/151.hea'\n",
      "Skipping 152 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/152.hea'\n",
      "Skipping 153 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/153.hea'\n",
      "Skipping 154 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/154.hea'\n",
      "Skipping 155 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/155.hea'\n",
      "Skipping 156 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/156.hea'\n",
      "Skipping 157 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/157.hea'\n",
      "Skipping 158 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/158.hea'\n",
      "Skipping 159 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/159.hea'\n",
      "Skipping 160 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/160.hea'\n",
      "Skipping 161 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/161.hea'\n",
      "Skipping 162 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/162.hea'\n",
      "Skipping 163 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/163.hea'\n",
      "Skipping 164 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/164.hea'\n",
      "Skipping 165 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/165.hea'\n",
      "Skipping 166 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/166.hea'\n",
      "Skipping 167 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/167.hea'\n",
      "Skipping 168 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/168.hea'\n",
      "Skipping 169 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/169.hea'\n",
      "Skipping 170 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/170.hea'\n",
      "Skipping 171 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/171.hea'\n",
      "Skipping 172 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/172.hea'\n",
      "Skipping 173 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/173.hea'\n",
      "Skipping 174 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/174.hea'\n",
      "Skipping 175 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/175.hea'\n",
      "Skipping 176 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/176.hea'\n",
      "Skipping 177 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/177.hea'\n",
      "Skipping 178 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/178.hea'\n",
      "Skipping 179 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/179.hea'\n",
      "Skipping 180 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/180.hea'\n",
      "Skipping 181 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/181.hea'\n",
      "Skipping 182 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/182.hea'\n",
      "Skipping 183 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/183.hea'\n",
      "Skipping 184 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/184.hea'\n",
      "Skipping 185 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/185.hea'\n",
      "Skipping 186 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/186.hea'\n",
      "Skipping 187 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/187.hea'\n",
      "Skipping 188 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/188.hea'\n",
      "Skipping 189 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/189.hea'\n",
      "Skipping 190 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/190.hea'\n",
      "Skipping 191 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/191.hea'\n",
      "Skipping 192 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/192.hea'\n",
      "Skipping 193 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/193.hea'\n",
      "Skipping 194 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/194.hea'\n",
      "Skipping 195 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/195.hea'\n",
      "Skipping 196 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/196.hea'\n",
      "Skipping 197 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/197.hea'\n",
      "Skipping 198 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/198.hea'\n",
      "Skipping 199 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/199.hea'\n",
      "Skipping 204 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/204.hea'\n",
      "Skipping 206 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/206.hea'\n",
      "Skipping 211 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/211.hea'\n",
      "Skipping 216 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/216.hea'\n",
      "Skipping 218 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/218.hea'\n",
      "Skipping 224 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/224.hea'\n",
      "Skipping 225 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/225.hea'\n",
      "Skipping 226 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/226.hea'\n",
      "Skipping 227 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/227.hea'\n",
      "Skipping 229 due to error: [Errno 2] No such file or directory: 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/mit-bih-arrhythmia-database-1.0.0/229.hea'\n",
      "Class distribution: {'NORM': 529, 'OTHER': 1500, 'RBBB': 107, 'AFIB': 102, 'LBBB': 51}\n",
      "Label mapping: {'1dAVb': 0, 'AFIB': 1, 'AFLT': 2, 'LBBB': 3, 'NORM': 4, 'OTHER': 5, 'RBBB': 6}\n",
      "Shape of X: (2289, 12, 1000, 1)\n",
      "Label mapping: {'1dAVb': 0, 'AFIB': 1, 'AFLT': 2, 'LBBB': 3, 'NORM': 4, 'OTHER': 5, 'RBBB': 6}\n"
     ]
    }
   ],
   "source": [
    "# records = ['100', '101', '102', '103', '104'] \n",
    "records = [f\"{i:03d}\" for i in range(100, 234)]  # Adjust as needed\n",
    "X, y, label_map = prepare_mitbih_dataset(mitbih_path, records)\n",
    "print(\"Shape of X:\", X.shape)  # Expected: (N, 12, 1000, 1)\n",
    "print(\"Label mapping:\", label_map)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "aba07dd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping 100 due to error: [Errno 2] No such file or directory: 'c:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ECG-classification-main/path/to/mitbih/100.atr'\n",
      "Skipping 101 due to error: [Errno 2] No such file or directory: 'c:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ECG-classification-main/path/to/mitbih/101.atr'\n",
      "Skipping 102 due to error: [Errno 2] No such file or directory: 'c:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ECG-classification-main/path/to/mitbih/102.atr'\n",
      "Skipping 103 due to error: [Errno 2] No such file or directory: 'c:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ECG-classification-main/path/to/mitbih/103.atr'\n",
      "Skipping 104 due to error: [Errno 2] No such file or directory: 'c:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ECG-classification-main/path/to/mitbih/104.atr'\n",
      "Skipping 105 due to error: [Errno 2] No such file or directory: 'c:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ECG-classification-main/path/to/mitbih/105.atr'\n",
      "1dAVb cases found: Counter()\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "import wfdb\n",
    "import os\n",
    "\n",
    "def check_1dAVb_presence(data_dir, record_list):\n",
    "    counter = Counter()\n",
    "    for rec in record_list:\n",
    "        try:\n",
    "            annotation = wfdb.rdann(os.path.join(data_dir, rec), 'atr')\n",
    "            for note in annotation.aux_note:\n",
    "                if '1AV' in note:\n",
    "                    counter[rec] += 1\n",
    "        except Exception as e:\n",
    "            print(f\"Skipping {rec} due to error: {e}\")\n",
    "    return counter\n",
    "\n",
    "# Example usage\n",
    "records = ['100', '101', '102', '103', '104', '105']  # or all 48\n",
    "avb_cases = check_1dAVb_presence('path/to/mitbih', records)\n",
    "print(\"1dAVb cases found:\", avb_cases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "592c02a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_ptbxl shape: (24156,)\n",
      "First few entries in y_ptbxl: ['AFIB' 'NORM' 'OTHER' 'OTHER' 'OTHER']\n",
      "x_1davb shape: (1, 12, 1000, 1)\n",
      "x_aflt shape: (1, 12, 1000, 1)\n",
      "Original X shape: (2289, 12, 1000, 1)\n",
      "[5 5 5 5 5]\n",
      "y_1davb_add: [0]\n",
      "(2289,)\n",
      "✅ New test set shape: (2291, 12, 1000, 1) (2291,)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Load the PTB-XL balanced training set\n",
    "train_data = np.load('X_train_balanced.npz', allow_pickle=True)\n",
    "y_data = np.load('y_train_balanced.npz', allow_pickle=True)\n",
    "X_ptbxl = train_data['X']\n",
    "y_ptbxl = y_data['y']\n",
    "\n",
    "# Your label order: ['1dAVb', 'AFIB', 'AFLT', 'LBBB', 'NORM', 'OTHER', 'RBBB']\n",
    "label_order = ['1dAVb', 'AFIB', 'AFLT', 'LBBB', 'NORM', 'OTHER', 'RBBB']\n",
    "\n",
    "print(\"y_ptbxl shape:\", y_ptbxl.shape)\n",
    "print(\"First few entries in y_ptbxl:\", y_ptbxl[:5])\n",
    "\n",
    "# Find an index for each of the missing classes\n",
    "index_1davb = np.where(y_ptbxl == '1dAVb')[0][0]\n",
    "index_aflt = np.where(y_ptbxl == 'AFLT')[0][0]\n",
    "\n",
    "int_label_1davb = label_order.index('1dAVb') # This will be 0\n",
    "int_label_aflt = label_order.index('AFLT')   # This will be 2\n",
    "# Convert the string labels for the added examples to their integer mapping\n",
    "y_1davb_add = np.array([label_order.index('1dAVb')]) # This will be [0]\n",
    "y_aflt_add = np.array([label_order.index('AFLT')])   # This will be [2]\n",
    "\n",
    "x_1davb = X_ptbxl[index_1davb]\n",
    "x_aflt = X_ptbxl[index_aflt]\n",
    "\n",
    "# Ensure the new samples have shape (1, 12, 1000, 1)\n",
    "x_1davb = x_1davb.reshape(12, 1000, 1)  # Reshape to (12, 1000, 1)\n",
    "x_aflt = x_aflt.reshape(12, 1000, 1)\n",
    "# Expand dimensions to add batch size\n",
    "x_1davb = np.expand_dims(x_1davb, axis=0)  # Add batch dimension\n",
    "x_aflt = np.expand_dims(x_aflt, axis=0)\n",
    "\n",
    "# Double-check all shapes are 4D\n",
    "print(\"x_1davb shape:\", x_1davb.shape)\n",
    "print(\"x_aflt shape:\", x_aflt.shape)\n",
    "print(\"Original X shape:\", X.shape)\n",
    "\n",
    "print(y[:5])  # Check first few labels\n",
    "print(\"y_1davb_add:\", y_1davb_add)\n",
    "print(y.shape)\n",
    "\n",
    "X = np.concatenate([X, x_1davb, x_aflt], axis=0)\n",
    "y = np.concatenate([y, y_1davb_add, y_aflt_add], axis=0)\n",
    "\n",
    "print(\"✅ New test set shape:\", X.shape, y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "f677fd14",
   "metadata": {},
   "outputs": [],
   "source": [
    "#save the X and y arrays to .npz files\n",
    "output_filename = 'mitbih_data.npz'\n",
    "np.savez(output_filename, X=X, y=y, label_map=label_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "c45ef67c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import wfdb\n",
    "from sklearn.preprocessing import MultiLabelBinarizer\n",
    "from scipy.signal import resample\n",
    "from collections import defaultdict\n",
    "\n",
    "def extract_diagnosis_from_hea(file_path):\n",
    "    \"\"\"\n",
    "    Extract diagnosis codes from a .hea file.\n",
    "    \n",
    "    Args:\n",
    "        file_path (str): Path to the .hea file.\n",
    "    \n",
    "    Returns:\n",
    "        list of str: List of diagnosis codes (as strings).\n",
    "    \"\"\"\n",
    "    with open(file_path, 'r') as f:\n",
    "        for line in f:\n",
    "            if line.startswith('#Dx:'):\n",
    "                # Get the part after \"#Dx:\"\n",
    "                dx_line = line.strip().split(':', 1)[1]\n",
    "                # Diagnosis codes can be comma-separated\n",
    "                dx_codes = [code.strip() for code in dx_line.split(',')]\n",
    "                return dx_codes\n",
    "    return []\n",
    "\n",
    "\n",
    "\n",
    "def prepare_ecg_arrhythmia_dataset_debug(data_dir, label_mapping, desired_labels, max_records=None):\n",
    "    import os\n",
    "    import wfdb\n",
    "    import numpy as np\n",
    "    from scipy.signal import resample\n",
    "    from sklearn.preprocessing import MultiLabelBinarizer\n",
    "    from collections import defaultdict\n",
    "    from glob import glob\n",
    "\n",
    "    # use glob to find all .hea files in the directory and its subdirectories\n",
    "    if not os.path.exists(data_dir):\n",
    "        print(f\"Directory {data_dir} does not exist.\")\n",
    "        return None, None\n",
    "    record_files = glob(data_dir + '/**/*.hea', recursive=True)\n",
    "    \n",
    "    # record_files = Path(data_dir).rglob('*.hea')\n",
    "        \n",
    "\n",
    "\n",
    "    # record_files = [f for f in os.listdir(data_dir) if f.endswith('.hea')]\n",
    "    print(f\"Found {len(record_files)} records in {data_dir}\")\n",
    "    if max_records:\n",
    "        record_files = record_files[:max_records]\n",
    "\n",
    "    X = []\n",
    "    y = []\n",
    "    class_counts = defaultdict(int)\n",
    "\n",
    "    total_checked = 0\n",
    "\n",
    "    for file in record_files:\n",
    "        record_name = file[:-4]\n",
    "        total_checked += 1\n",
    "\n",
    "        try:\n",
    "            record = wfdb.rdrecord(os.path.join(data_dir, record_name))\n",
    "            signal = record.p_signal.T  # shape (12, time)\n",
    "\n",
    "            if signal.shape[1] < 5000:\n",
    "                print(f\"Skipping {record_name} (too short: {signal.shape[1]})\")\n",
    "                continue\n",
    "\n",
    "            resampled = resample(signal, 1000, axis=1)\n",
    "\n",
    "            # Attempt to get annotations\n",
    "            # annotations are found in the .hea files\n",
    "            # look for '#Dx' line in the .hea file and parse it\n",
    "            annotation = extract_diagnosis_from_hea(os.path.join(data_dir, file))\n",
    "            # print(f\"Annotations for {record_name}: {annotation}\")\n",
    "            \n",
    "            labels = set()\n",
    "            \n",
    "            if annotation:\n",
    "                for note in annotation:\n",
    "                    \n",
    "                    mapped = label_mapping.get(note)\n",
    "                    if mapped:\n",
    "                        labels.add(mapped)\n",
    "            else:\n",
    "                print(f\"Missing annotation for {record_name}\")\n",
    "\n",
    "            # Fallback to OTHER if nothing found\n",
    "            if not labels:\n",
    "                labels = {'OTHER'}\n",
    "            else:\n",
    "                labels = {lbl if lbl in desired_labels else 'OTHER' for lbl in labels}\n",
    "                if 'OTHER' in labels and len(labels) > 1:\n",
    "                    labels.discard('OTHER')\n",
    "\n",
    "            label = list(labels)[0]\n",
    "            class_counts[label] += 1\n",
    "\n",
    "            X.append(resampled)\n",
    "            y.append(label)\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Failed on {record_name}: {e}\")\n",
    "\n",
    "    print(f\"\\n🔍 Checked {total_checked} records, loaded {len(X)}\")\n",
    "    for k, v in class_counts.items():\n",
    "        print(f\"  {k}: {v}\")\n",
    "\n",
    "    if len(X) == 0:\n",
    "        print(\"⚠️ No data loaded. Double-check annotation file format and mapping keys.\")\n",
    "        return None, None\n",
    "\n",
    "    X = np.array(X)[..., np.newaxis]\n",
    "    mlb = MultiLabelBinarizer(classes=desired_labels)\n",
    "    y = mlb.fit_transform([[lbl] for lbl in y])\n",
    "\n",
    "    return X, y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "90bd1797",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 45152 records in C:\\Users\\mstew\\OneDrive\\Curtin Univeristy\\COMP6011\\Task 3\\a-large-scale-12-lead-electrocardiogram-database-for-arrhythmia-study-1.0.0\\WFDBRecords\\\n",
      "Failed on C:\\Users\\mstew\\OneDrive\\Curtin Univeristy\\COMP6011\\Task 3\\a-large-scale-12-lead-electrocardiogram-database-for-arrhythmia-study-1.0.0\\WFDBRecords\\01\\019\\JS01052: time data '/' does not match format '%d/%m/%Y'\n",
      "\n",
      "🔍 Checked 5000 records, loaded 4999\n",
      "  RBBB: 177\n",
      "  OTHER: 2650\n",
      "  AFLT: 239\n",
      "  AFIB: 802\n",
      "  NORM: 944\n",
      "  1dAVb: 107\n",
      "  LBBB: 80\n"
     ]
    }
   ],
   "source": [
    "label_mapping = {\n",
    "    '270492004': '1dAVb',\n",
    "    '164889003': 'AFIB',\n",
    "    '164890007': 'AFLT',\n",
    "    '164909002': 'LBBB',\n",
    "    '426783006': 'NORM',\n",
    "    '59118001': 'RBBB',\n",
    "    # Unmapped labels will be classified as OTHER\n",
    "}\n",
    "\n",
    "desired_labels = ['1dAVb', 'AFIB', 'AFLT', 'LBBB', 'NORM', 'OTHER', 'RBBB']\n",
    "\n",
    "X_ecg, y_ecg = prepare_ecg_arrhythmia_dataset_debug(\n",
    "    data_dir='C:\\\\Users\\\\mstew\\\\OneDrive\\\\Curtin Univeristy\\\\COMP6011\\\\Task 3\\\\a-large-scale-12-lead-electrocardiogram-database-for-arrhythmia-study-1.0.0\\WFDBRecords\\\\',\n",
    "    label_mapping=label_mapping,\n",
    "    desired_labels=desired_labels,\n",
    "    max_records=5000\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "80866c6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3499, 12, 1000, 1)\n",
      "(3499, 7)\n",
      "Shapes before removing OTHER: (3499, 12, 1000, 1) (3499, 7)\n",
      "Shapes after removing OTHER: (3499, 12, 1000, 1) (3499, 7)\n",
      "y_ecg_single shape: (3499,)\n",
      "First few labels: [6 5 2 1 4]\n"
     ]
    }
   ],
   "source": [
    "print(X_ecg.shape)\n",
    "print(y_ecg.shape)\n",
    "\n",
    "# remove 1500 samples from OTHER class\n",
    "if X_ecg is not None and y_ecg is not None:\n",
    "    print(\"Shapes before removing OTHER:\", X_ecg.shape, y_ecg.shape)\n",
    "    \n",
    "    # Find indices of OTHER class\n",
    "    other_indices = np.where(y_ecg[:, 5] == 1)[0]  # Assuming 'OTHER' is at index 5\n",
    "    if len(other_indices) > 1500:\n",
    "        indices_to_remove = np.random.choice(other_indices, size=1500, replace=False)\n",
    "        mask = np.ones(len(y_ecg), dtype=bool)\n",
    "        mask[indices_to_remove] = False\n",
    "        X_ecg = X_ecg[mask]\n",
    "        y_ecg = y_ecg[mask]\n",
    "    \n",
    "    print(\"Shapes after removing OTHER:\", X_ecg.shape, y_ecg.shape)\n",
    "\n",
    "# change the 1 hot encoding of y_ecg to a single label\n",
    "y_ecg_single = np.argmax(y_ecg, axis=1)  # Convert to single label per sample\n",
    "\n",
    "print(\"y_ecg_single shape:\", y_ecg_single.shape)\n",
    "print(\"First few labels:\", y_ecg_single[:5])\n",
    "\n",
    "# save the X and y arrays to .npz files\n",
    "output_filename = 'ecg_arrhythmia_data.npz'\n",
    "np.savez(output_filename, X=X_ecg, y=y_ecg_single, label_map=desired_labels)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ECG_39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
